import torch
from torch import nn
import torch.nn.functional as F


class TwoTower(nn.Module):
    def __init__(self, n_user,
                 n_age,
                 n_job,
                 n_zip,
                 n_gender,
                 n_movie,
                 n_genre,
                 n_name_token_dim,
                 embed_dim,
                 hidden_size,
                 ):
        super().__init__()
        assert embed_dim % 2 == 0

        self.user_embeds = nn.ModuleList([
            nn.Embedding(n_user, embed_dim),
            nn.Embedding(n_gender, embed_dim // 2),
            nn.Embedding(n_age, embed_dim // 2),
            nn.Embedding(n_job, embed_dim // 2),
            nn.Embedding(n_zip, embed_dim // 2),
        ])

        self.user_linear = nn.ModuleList([
            nn.Linear(embed_dim, embed_dim),
            nn.Linear(embed_dim // 2, embed_dim // 2),
            nn.Linear(embed_dim // 2, embed_dim // 2),
            nn.Linear(embed_dim // 2, embed_dim // 2),
            nn.Linear(embed_dim // 2, embed_dim // 2)

        ])

        user_in_dim = embed_dim + embed_dim // 2 * 4
        self.user_cat_fc = nn.Linear(user_in_dim, hidden_size)

        self.movie_embedding = nn.ModuleList([
            nn.Embedding(n_movie, embed_dim),
            None,
            nn.Embedding(n_genre, embed_dim // 2, padding_idx=0)
        ])

        self.movie_fc = nn.ModuleList([
            nn.Linear(embed_dim, embed_dim),
            nn.Linear(n_name_token_dim, embed_dim),
            nn.Linear(embed_dim // 2, embed_dim // 2),
        ])

        movie_in_dim = embed_dim + embed_dim + embed_dim // 2
        self.movie_cat_fc = nn.Linear(movie_in_dim, hidden_size)

        self.fc = nn.Linear(hidden_size * 2, 1)
        self.normal_output = NormalizeToRange()

    def forward(self, user_input, movie_input):
        for i, (k, v) in enumerate(user_input.items()):
            user_input[k] = self.user_embeds[i](v)
            user_input[k] = self.user_linear[i](user_input[k])

        user_cat = torch.cat(list(user_input.values()), dim=-1)

        user_feature = F.tanh(self.user_cat_fc(user_cat))

        for i, (k, v) in enumerate(movie_input.items()):
            if i == 1:  # 当输出是movie_name时
                movie_input[k] = self.movie_fc[i](v)
            elif i == 2: # 当输出是movie_genre时
                tmp = self.movie_embedding[i](v)

                mean_genre_embedding = torch.mean(tmp, dim=1)

                movie_input[k] = self.movie_fc[i](mean_genre_embedding)
            else:


                movie_input[k] = self.movie_embedding[i](v)
                movie_input[k] = self.movie_fc[i](movie_input[k])

        movie_cat = torch.cat(list(movie_input.values()), dim=-1)
        movie_feature = F.tanh(self.movie_cat_fc(movie_cat))
        features_cat = torch.cat([user_feature, movie_feature], dim=-1)
        output = self.fc(features_cat)
        output = self.normal_output(output)
        return output, user_feature, movie_feature


# 归一化电影评分
class NormalizeToRange(nn.Module):
    def __init__(self, min_val=1, max_val=5):
        super(NormalizeToRange, self).__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        sigmoid_output = self.sigmoid(x)
        return sigmoid_output * (self.max_val - self.min_val) + self.min_val






